Repository navigation
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a flexible benchmarking framework for Host-to-Device (H2D) and Device-to-Host (D2H) transfers, supporting multiple strategies such as simple, pipelined, and pinned memory transfers. It includes updated configuration files for comprehensive experiments and new scripts to automate execution with optional NUMA bindings. The review comments highlight critical improvements for multi-device data distribution, robustness in variance calculations, and safer execution of shell scripts when dependencies like numactl are missing.
I am having trouble creating individual review comments. Click here to see my feedback.
Ironwood/src/benchmark_host_device.py (126-132)
In a multi-device benchmark configuration, falling back to transferring the entire array to self.target_devices[0] when self.chunks_per_dev == 1 ignores the other devices. This makes the benchmark results incorrect and inconsistent for small data sizes. Instead, we should split the data across all target devices (1 chunk per device) to ensure a fair and correct multi-device transfer benchmark.
else:
chunks = np.array_split(host_data, self.num_devices, axis=0)
tensors_on_device = []
t0 = time.perf_counter()
for idx, chunk in enumerate(chunks):
dev = self.target_devices[idx]
tensors_on_device.append(jax.device_put(chunk, dev))
for device_tensor in tensors_on_device:
device_tensor.block_until_ready()
t1 = time.perf_counter()
self.h2d_perf.append((t1 - t0) * 1000)
del chunks
return tensors_on_device
Ironwood/src/benchmark_host_device.py (101-105)
Since we are updating the fallback behavior to split across all devices instead of transferring only to a single device, we should update the warning message to accurately reflect this fallback strategy.
if self.chunks_per_dev == 1:
print(
"Warning: Data size is not larger than target chunk size, "
"falling back to 1 chunk per device."
)
Ironwood/src/benchmark_host_device.py (317-320)
Calculating sample variance with ddof=1 when len(bw_array) <= 1 will result in a division by zero, causing a RuntimeWarning and returning nan. We should check the length of bw_array before calculating the variance to make the code more robust.
bw_array = np.array(bw_list)
if len(bw_array) > 1:
sample_variance = np.var(bw_array, ddof=1)
else:
sample_variance = 0.0
Ironwood/guides/host_device_benchmark/run_comprehensive.sh (12-13)
If numactl is not installed on the system, the script will fail and exit immediately due to set -e. It is safer to check if numactl is available before executing it, falling back to running without it if necessary.
echo "Running 8-device benchmark with numactl interleaving..."
if command -v numactl &> /dev/null; then
numactl --interleave=all python3 src/run_benchmark.py --config configs/host_device/comprehensive_8dev_experiments.yaml
else
echo "Warning: numactl not found. Running 8-device benchmark without interleaving."
python3 src/run_benchmark.py --config configs/host_device/comprehensive_8dev_experiments.yaml
fi
linamy85
left a comment
linamy85
left a comment
There was a problem hiding this comment.
Thanks Simran! The pinned H2D and D2H timing both look correct: allocation is outside timing, all 8 shards transfer concurrently, and block_until_ready is inside the timer. Main asks: assert the device count and fix the jax/libtpu pin. Since we're moving this to main soon, I've skipped design comments.
- Add assertions in SimpleTransfer, PipelinedTransfer, and PinnedMemoryTransfer to guarantee len(target_devices) == num_devices, preventing silent execution with fewer devices than requested. - Update jax[tpu] and jaxlib to 0.12.0 and libtpu to 0.0.50 to resolve the dependency conflict with libtpu 0.0.48 and fix the D2H regression.
Summary
Ports the pinned memory Host-to-Device (
H2D) and Device-to-Host (D2H) transfer benchmark (PinnedMemoryTransfer,SimpleTransfer, andPipelinedTransfer) and associated configurations/scripts from thetpu7x-h2dd2hbranch tochs, along withjax.local_devices()sharding for single-host/node-local execution andnumactlsupport inDockerfile.Code References / Source History
This PR ports the changes introduced on
tpu7x-h2dd2hacross the following commits:9b4e8de("Try pinned memory" by@linamy85): AddedTransferStrategyhierarchy (SimpleTransfer,PipelinedTransfer,PinnedMemoryTransferusingNamedSharding(..., memory_kind="pinned_host")for H2D and parallel shardjax.device_getfor D2H) inIronwood/src/benchmark_host_device.py.f924a7e("Optimize H2D/D2H transfer pipelines and add comprehensive benchmark configs" by@linamy85): UpdatedIronwood/src/benchmark_host_device.pyandIronwood/configs/host_device/host_device.yaml, and addedIronwood/configs/host_device/comprehensive_experiments.yamlandIronwood/configs/host_device/comprehensive_8dev_experiments.yaml.7439b2a("Add benchmark guide and run script" by@linamy85): UpdatedIronwood/scripts/run_host_device_benchmark.sh(TPU_PREMAPPED_BUFFER_SIZE=34359738368,--numactl_interleave), and addedIronwood/guides/host_device_benchmark/README.mdandIronwood/guides/host_device_benchmark/run_comprehensive.sh.38ec530("Allow sweeping dtype in host_device benchmarks" by@linamy85): Added configurabledtypesupport inIronwood/src/benchmark_host_device.py.Additional
chsChanges in This PRDockerfile:numactlsonumactl --cpunodebind=0 --membind=0andnumactl --interleave=allruns work inside the container.requirements.txt:jax[tpu]==0.11.1andjaxlib==0.11.1while keepinglibtpu==0.0.48(matchingmainpyproject.toml), avoiding thejax 0.11.2jax.device_getD2H and multi-device H2D performance regressions while retaininglibtpu 0.0.48forsingle_host_tpu_init.Tests: https://docs.google.com/spreadsheets/d/1VQqHS1UEvbgvyZ6ddnl6-Gy-79fwYNTkSSYopMsQnC8/edit?resourcekey=0-sFolkGOQhiSUVIthaPPszg&gid=920016885#gid=920016885